use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserTraffic {
    pub tx: u64,
    pub rx: u64,
    pub online: bool,
    pub connections: u32,
    pub last_seen: DateTime<Utc>,
}

pub struct TrafficStats {
    users: Arc<RwLock<HashMap<String, UserTraffic>>>,
}

impl TrafficStats {
    pub fn new() -> Self {
        Self {
            users: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn log_traffic(&self, user_id: &str, tx: u64, rx: u64) {
        let mut users = self.users.write().await;
        let entry = users.entry(user_id.to_string()).or_insert_with(|| UserTraffic {
            tx: 0,
            rx: 0,
            online: true,
            connections: 0,
            last_seen: Utc::now(),
        });

        entry.tx += tx;
        entry.rx += rx;
        entry.last_seen = Utc::now();
    }

    pub async fn set_online(&self, user_id: &str, online: bool) {
        let mut users = self.users.write().await;
        if let Some(entry) = users.get_mut(user_id) {
            entry.online = online;
            entry.last_seen = Utc::now();
            
            if online {
                entry.connections += 1;
            } else if entry.connections > 0 {
                entry.connections -= 1;
            }
        } else if online {
            users.insert(
                user_id.to_string(),
                UserTraffic {
                    tx: 0,
                    rx: 0,
                    online: true,
                    connections: 1,
                    last_seen: Utc::now(),
                },
            );
        }
    }

    pub async fn get_all_users(&self) -> HashMap<String, UserTraffic> {
        self.users.read().await.clone()
    }

    pub async fn get_user(&self, user_id: &str) -> Option<UserTraffic> {
        self.users.read().await.get(user_id).cloned()
    }

    pub async fn clear_stats(&self) {
        self.users.write().await.clear();
    }

    pub async fn get_total(&self) -> (u64, u64) {
        let users = self.users.read().await;
        let mut total_tx = 0u64;
        let mut total_rx = 0u64;

        for traffic in users.values() {
            total_tx += traffic.tx;
            total_rx += traffic.rx;
        }

        (total_tx, total_rx)
    }
}

impl Default for TrafficStats {
    fn default() -> Self {
        Self::new()
    }
}
